# -*- coding: utf-8 -*-
"""
Created on Mon Jan 20 23:15:34 2020
@author: ikaros
"""
'''
该程序用于预处理数据集，由于弃用了EAST的dataset.py的crop_img函数，
用此文件处理完后才可以用EAST的train.py运行，
该文件主要用于把图片变成256*256的尺寸，有两种模式，
模式一针对卡片式子的图片，由于只有一个目标，直接resize成256*256.
模式二针对一页的算式图片，由于一页的尺寸w/h约等于1/2，故resize成256*512，再切成上下两半分成两张256*256的图。

主要的处理对象为模式二，
原来的EAST只要是文字就会当成目标，因此crop后目标区域的坐标变换方法不是很恰当，
一个字被切成两半也会被当成目标，引入了很大的误差，
但我们对式子完整性要求很高，因为完整性影响着后续的文字识别功能，不完整定位后的式子会使识别产生很大的错误，
于是，我这里以0.85位阈值，处于分界线的目标，被切成两半后如果还保留着超过0.85的部分就会被当成目标识别，
小于0.85的就认为该式子不完整，不能被当作正样本来训练。
'''

import numpy as np
import cv2
from PIL import Image
import math
import os #路径有中文时可能会报一些奇怪的错


o_img_path='./o_img/'                   #原始图片地址
save_img_dir = './compress_img/'          #新图片保存地址
o_label_dir = './o_label/'              #原始标签地址
save_label_dir1 = './compress_label/'      #新标签保存地址
save_label_dir2 = './compress_fix_label/'
checkdir = './check/'


def extract_vertices(lines):
	'''extract vertices info from txt lines
	Input:
		lines   : 标签文件里的所有内容
	Output:
		vertices: vertices of text regions <numpy.ndarray, (n,8)>
		labels  : 1->valid, 0->ignore, <numpy.ndarray, (n,)>
	'''
	labels = []
	vertices = []
	for line in lines:
		vertices.append(list(map(int,list(map(float,line.rstrip('\n').rstrip().lstrip('\ufeff').split(',')[:8])))))
		label = 0 if '###' in line else 1
		labels.append(label)
	return np.array(vertices), np.array(labels)

def mode_decision(img_np,lines):
    '''
    判断图片是哪种模式：
    模式一：图片中目标宽度大于等于图像宽度的一半。------>用于处理卡片
    模式二：图片中目标宽度小于图像宽度的一半。   ------->用于处理算式
    '''
    img_w = img_np.shape[1]
    vertices, label =extract_vertices(lines)
    x_min = 0
    x_max = 0
    for v1 in vertices:
        if v1 is None:
            print('None')
            break
        x1 = round(v1[0])
        x2 = round(v1[2])
        x3 = round(v1[4])
        x4 = round(v1[6])
        
        x = [x1,x2,x3,x4]
        x_min = min(x)
        x_max = max(x)
        
        obj_w = x_max - x_min
        if not(obj_w < (img_w*0.5)):
            mode = 1
            return mode
    mode = 2
    return mode

def mode1_resize(img,lines,save_imgname,save_labelname):
    '''
    根据模式对图片进行不同尺寸的resize。
    模式一：256*256
    模式二：256*512，然后再次切成两半256*256
    '''
    img_w = img.shape[1]
    img_h = img.shape[0]
    vertices, label =extract_vertices(lines)
    
    scale_w = 256 / img_w
    scale_h = 256 / img_h
    
    img2 = cv2.resize(img,(256,256))
    cv2.imwrite(save_imgname,img2)
    
    newlines = []
    for v1 in vertices:
        if v1 is None:
            print('None')
            break
        x1 = str(int(round(v1[0] * scale_w)))
        y1 = str(int(round(v1[1] * scale_h)))
        x2 = str(int(round(v1[2] * scale_w)))
        y2 = str(int(round(v1[3] * scale_h)))
        x3 = str(int(round(v1[4] * scale_w)))
        y3 = str(int(round(v1[5] * scale_h)))
        x4 = str(int(round(v1[6] * scale_w)))
        y4 = str(int(round(v1[7] * scale_h)))
        
        newline = x1 + ',' + y1 + ',' + x2 + ',' + y2 + ','+ x3 + ','+ y3 + ','+ x4 + ','+ y4 + ',' + 'A'+'\n'
        newlines.append(newline)
        
        with open(save_labelname, "a") as f1:         #把坐标和标签内容写入txt文件
            for ns in newlines:
                f1.write(ns)

def mode2_resize(img,lines,img_name,save_imgdir,save_labeldir):
    img_w = img.shape[1]
    img_h = img.shape[0]
    vertices, label =extract_vertices(lines)
    scale_w = 256 / img_w
    scale_h = 512 / img_h
    
    img2 = cv2.resize(img,(256,512))
    #img2 = cv2.resize(img,(512,256))
    
    '''去掉被边界截到的目标'''
    uponlines = []
    belowlines = []
    flag1 = 0
#    yuzhi = 6                  #根据压缩图像后目标框的高度来取舍
    for v1 in vertices:
#        if v1 is None:
#            print('None')
#            break
        x1 = int(round(v1[0] * scale_w))
        y1 = int(round(v1[1] * scale_h))
        x2 = int(round(v1[2] * scale_w))
        y2 = int(round(v1[3] * scale_h))
        x3 = int(round(v1[4] * scale_w))
        y3 = int(round(v1[5] * scale_h))
        x4 = int(round(v1[6] * scale_w))
        y4 = int(round(v1[7] * scale_h))
        
        y = [y1,y2,y3,y4]
        y_min = min(y)
        y_max = max(y)
        h_min = y_max - y_min
#        if h_min < 5:
#            flag1 = 0
#            continue
        
#        if (y_min < 256) and (y_max >256):
#            continue
        
        x1 = str(x1)
        x2 = str(x2)
        x3 = str(x3)
        x4 = str(x4)
        if (y_max <= 255):
            flag1 = 1
        elif (y_min < 255) and (y_max >255):
            if (y_min + h_min*0.85) <= 255:
                flag1 = 1
            elif (y_max - h_min*0.85) >= 255:
                flag1 = 2
#            if ((y_max-256)<yuzhi):
#                flag1 = 1
#            elif ((y_max-256)>1) and ((256-y_min) < yuzhi):
#                flag1 = 2
            else:
                flag1 = 0
        elif (y_min >= 255):
            flag1 = 2
        else:
            flag1 = 0
            
        if flag1 == 1:
            y1 = str(y1)
            y2 = str(y2)
            y3 = str(y3)
            y4 = str(y4)
            upline = x1 + ',' + y1 + ',' + x2 + ',' + y2 + ','+ x3 + ','+ y3 + ','+ x4 + ','+ y4 + ',' + 'A'+'\n'
            uponlines.append(upline)
            flag1 = 0
        elif flag1 == 2:
            y1 = str(int(y1-256))
            y2 = str(int(y2-256))
            y3 = str(int(y3-256))
            y4 = str(int(y4-256))
            belowline = x1 + ',' + y1 + ',' + x2 + ',' + y2 + ','+ x3 + ','+ y3 + ','+ x4 + ','+ y4 + ',' + 'A'+'\n'     
            belowlines.append(belowline)
            flag1 = 0
        else:
            pass
            
    
    np_up = np.array(uponlines)
    np_below = np.array(belowlines)
    if np_up.size != 0:
        uplabel = save_labeldir + img_name.replace('.jpg','(upon).txt')
        with open(uplabel, "a") as f1:         #把坐标和标签内容写入txt文件
            for ns in uponlines:
                f1.write(ns)
    if np_below.size != 0:
        belowlabel = save_labeldir + img_name.replace('.jpg','(below).txt')
        with open(belowlabel, "a") as f2:         #把坐标和标签内容写入txt文件
            for ns in belowlines:
                f2.write(ns)

    #img_upon = cv2.resize((img2[0:256,0:bound_max]),(256,256))
    #img_below = cv2.resize((img2[0:256,bound_max:512]),(256,256))
    
    img_upon = img2[0:256,0:256]                #注意先高h后宽w
    img_below = img2[256:512,0:256]
    img_below = cv2.resize(img_below,(256,256))
    
    upon_name = save_imgdir + img_name.replace('.jpg','(upon).jpg')
    below_name = save_imgdir + img_name.replace('.jpg','(below).jpg')
    
    cv2.imwrite(upon_name,img_upon)
    cv2.imwrite(below_name,img_below)
    
    return upon_name,below_name

def verticle_repair(labeldir,label_savedir2):
    '''
    upon_names:  类型为list，用mode2_resize处理后的所有upon_name存到这个列表里
    below_names: 类型为list，用mode2_resize处理后的所有below_name存到这个列表里
    savedir2:    保存新的修正坐标后的label文件，地址要与upon_name、below_name的根目录不同
    '''
    labellist = os.listdir(labeldir)
    for labelname in labellist:
        read_label = labeldir + labelname
        with open(read_label, "r") as f1:            #读取标签文件的坐标和标签
            lines = f1.readlines()
        vertices, label3 = extract_vertices(lines)  #提取坐标
        newlines=[]
        for v1 in vertices:                         #判断坐标y是否超出图片范围，超出就修正回来
            if v1 is None:
                print('None')
                break
            x1 = str(int(round(v1[0])))
            y1 = round(v1[1])
            x2 = str(int(round(v1[2])))
            y2 = round(v1[3])
            x3 = str(int(round(v1[4])))
            y3 = round(v1[5])
            x4 = str(int(round(v1[6])))
            y4 = round(v1[7])
            
            y = [y1,y2,y3,y4]
            y_min = min(y)
            y_max = max(y)
            h_min = y_max - y_min
            if h_min < 5:
                continue
            
            if y1 < 0:      
                y1 = 0
            elif y1 > 255:
                y1 = 255
            y1 = str(y1)
            
            if y2 < 0:
                y2 = 0
            elif y2 > 255:
                y2 = 255
            y2 = str(y2)
                
            if y3 < 0:
                y3 = 0
            elif y3 > 255:
                y3 = 255
            y3 = str(y3)
                
            if y4 < 0:
                y4 = 0
            elif y4 > 255:
                y4 = 255
            y4 = str(y4)
            
            newline = x1 + ',' + y1 + ',' + x2 + ',' + y2 + ','+ x3 + ','+ y3 + ','+ x4 + ','+ y4 + ',' + 'A'+'\n'
            newlines.append(newline)
        
        new_labelname = label_savedir2 + labelname
        with open(new_labelname, "a") as f2:         #把坐标和标签内容写入txt文件
            for ns in newlines:
                f2.write(ns)


def check_verticles(imgdir,labeldir,out_dir):
    '''
    旋转完后图片后，
    用来检查坐标是否超出图片的大小范围，
    以及检查是否有的图片丢失标签文件。
    '''
    #若标签数多于图片数，
    #可以遍历标签名，把标签名转化图片名，
    #使用if not (img_names == imgname).any() 进行判断
    ###但一般都是图片比标签多，因为人工标签遗漏了部分图片。
    '''
    使用方法：
    注释掉旋转图片代码段，反注释最后一行检查函数语句。
    '''
    for root,dirs,files in os.walk(imgdir): #读取所有图片的地址和名字 #路径不能有中文
        pass
    label_names = np.array(os.listdir(labeldir))         #获取所有标签文件名
    num=0
    hmaxs =[]
    for file in files:
        imgname = imgdir + file
        img = cv2.imread(imgname)
        w = np.array(img).shape[1]                       #宽高用来判断坐标是否超出图片范围，从而导致训练出错
        h = np.array(img).shape[0]
        num += 1                                         #用来计数
        #labelname = 'gt_' + file.split('.')[0] + '.txt'  #把图片名转换成标签名
        labelname = file.split('.')[0] + '.txt'  #把图片名转换成标签名
        #print(labelname)
        
        '''#判断标签数是否少于图片数'''
        if not (label_names == labelname).any():         
            print('no label-->labelname:',labelname,' False')
        
        label2 = labeldir + labelname
        if (label_names == labelname).any():
            with open(label2, "r") as f1:                #读取标签文件的坐标和标签
                lines = f1.readlines()
            vertices, label3 = extract_vertices(lines)   #提取坐标
            lines2 = np.array(lines)
            if lines2.shape[0] == 0:                     #检查标签文件内容是否为空
                print('labelname--','shape wrong:',labelname,':',lines2.shape[0])
            '''把坐标取整'''
            for v1 in vertices:
                if v1 is None:
                    print('None')
                    break
                x1 = round(v1[0])
                y1 = round(v1[1])
                x2 = round(v1[2])
                y2 = round(v1[3])
                x3 = round(v1[4])
                y3 = round(v1[5])
                x4 = round(v1[6])
                y4 = round(v1[7])
                
                y = [y1,y2,y3,y4]
                ymin = min(y)
                ymax = max(y)
                hmax = ymax - ymin
                hmaxs.append(hmax)
                
                if x1<0 or x2<0 or x3<0 or x4<0 or y1<0 or y2<0 or y3<0 or y4<0:
                    print('v_wrong labelname:',labelname,' False')
                if x1>w or x2>w or x3>w or x4>w:
                    print('v_wrong labelname:',labelname,' False')
                if y1>h or y2>h or y3>h or y4>h:
                    print('v_wrong labelname:',labelname,' False')
                
                '''用直线连接 绘制四边形'''
                '''【注意box的格式】'''
                box = np.array([[x1,y1],[x2,y2],[x3,y3],[x4,y4]],np.int32)
                box = box.reshape(-1,1,2)
                cv2.polylines(img, [box], True,color=(0, 255, 0),thickness=2)
            #img = cv2.resize(img,(500,500))
            outname = out_dir + 'gt_' + file.split('.')[0] + '.jpg'
            cv2.imwrite(outname,img)
    print('h:',min(hmaxs))
        
def remove_wrong_img(labeldir,imgdir):
    '''删除没有目标的图片'''
    labelnames = os.listdir(labeldir)
    imgnames = os.listdir(imgdir)
    
    for imgname in imgnames:
        labelname = imgname.replace('.jpg','.txt')
        if not labelname in labelnames:
            img = imgdir + imgname
            os.remove(img)

def run_resize(o_imgdir,o_labeldir,save_imgdir,save_labeldir1):
    imgslist = os.listdir(o_imgdir)
    
    for imgname in imgslist:
        imgdir = o_imgdir + imgname
        img = cv2.imread(imgdir)
        img_np = np.array(img)
        
        labelname = o_labeldir + 'gt_' + imgname.replace('.jpg','.txt')
        save_labelname = save_labeldir1 + imgname.replace('.jpg','.txt')
        with open(labelname, "r") as f1:                #读取标签文件的坐标和标签
            lines = f1.readlines()
        
        mode = mode_decision(img_np,lines)
        if mode == 1:
            save_imgname = save_imgdir + imgname
            mode1_resize(img,lines,save_imgname,save_labelname)
        elif mode == 2:
            upon_name,below_name = mode2_resize(img,lines,imgname,save_imgdir,save_labeldir1)
        else:
            pass


run_resize(o_img_path,o_label_dir,save_img_dir,save_label_dir1)
            
verticle_repair(save_label_dir1,save_label_dir2)

remove_wrong_img(save_label_dir2,save_img_dir)



'''*******检查图片是否出错********'''
check_verticles(save_img_dir,save_label_dir2,checkdir)






























































































